#! usr/bin/env python
import cv2
import numpy as np
# import matplotlib.pyplot as plt

def make_coordinates(image, line_parameters):
    try:
        slope, intercept = line_parameters
    except TypeError:
        slope, intercept = 0.001, 0

    y1 = int(image.shape[0])
    y2 = int(y1 * (3/5))
    x1 = int((y1 - intercept) / slope)
    x2 = int((y2 - intercept) / slope)
    return [x1, y1, x2, y2]

def average_slope_intercept(image, lines):
    left_fit = []
    right_fit= []

    if lines is None:
        return None

    for line in lines:
        x1, y1, x2, y2 = line.reshape(4)
        parameters = np.polyfit((x1, x2), (y1, y2), 1) # Fit a polynomial of degree 1 and return the m and k
        slope = parameters[0]
        intercept = parameters[1]

        if slope < 0:
            left_fit.append((slope, intercept)) # The y axis is going downwards positively, so left lines actually result in negative slope
        else:
            right_fit.append((slope, intercept))

    # if len(left_fit) and len(right_fit):
    left_fit_average = np.average(left_fit, axis=0)
    right_fit_average = np.average(right_fit, axis=0)
    left_line = make_coordinates(image, left_fit_average)
    right_line = make_coordinates(image, right_fit_average)

    return [left_line, right_line]

def canny(image):
    grayed_img = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    blurred_img = cv2.GaussianBlur(grayed_img, (5, 5), 0)
    canny = cv2.Canny(blurred_img, 50, 150)
    return canny

def display_lines(image, lines):
    line_image = np.zeros_like(image)
    if lines is not None:
        for line in lines:
            # x1, y1, x2, y2 = line.reshape(4)
            x1, y1, x2, y2 = line
            cv2.line(line_image, (x1, y1), (x2, y2), (255, 0, 0), 10)
    return line_image

def region_of_interest(image):
    height= image.shape[0]
    polygons = np.array([[(200, height), (1100, height), (550, 250)]])
    mask = np.zeros_like(image)
    cv2.fillPoly(mask, polygons, 255)
    masked_img = cv2.bitwise_and(image, mask)
    return masked_img


cap = cv2.VideoCapture("../test2.mp4")
while(cap.isOpened()):

    ret, frame = cap.read()
    if ret:
        canny_image = canny(frame)
        cropped_img = region_of_interest(canny_image)

        lines = cv2.HoughLinesP(cropped_img, 2, np.pi / 180, 100,
                                np.array([]), minLineLength=40, maxLineGap=5)
        average_lines = average_slope_intercept(frame, lines)
        line_image = display_lines(frame, average_lines)
        combo_image = cv2.addWeighted(frame, 0.8, line_image, 1, 1)

        cv2.imshow("Result", combo_image)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    else:
        break

cap.release()
cv2.destroyAllWindows()
